# coding:utf-8
from image_info import ImageInfo
from position import Position
import os
import sys
from PIL import Image
import logging
import random
from file_record import FileRecord
import shutil


class ImageData(object):
    def __init__(self, info_file, pos_left_up, pos_right_down, path_prefix=None):
        self.data = []
        with open(info_file) as f:
            for line in f:
                line = line.split(' ')
                self.data.append(ImageInfo(os.path.join(path_prefix, line[0]),
                                           Position(float(line[1].split(',')[0]), float(line[1].split(',')[1])),
                                           Position(float(line[2].split(',')[0]), float(line[2].split(',')[1])),
                                           int(line[3]),
                                           int(line[4]),
                                           ))

        self.data.sort()

        """找标准的左上角"""
        min_dist = sys.maxint
        for i in range(len(self.data)):
            if self.data[i].left_up.dist(pos_left_up) < min_dist:
                min_dist = self.data[i].left_up.dist(pos_left_up)
                self._st_min_index = i

        """找标准的右下角"""
        min_dist = sys.maxint
        for i in range(len(self.data)):
            if self.data[i].left_up.dist(pos_right_down) < min_dist:
                min_dist = self.data[i].left_up.dist(pos_right_down)
                self._st_max_index = i

        assert(self._st_min_index < self._st_max_index)

        self._st_min_pos = self.data[self._st_min_index].left_up
        self._st_max_pos = self.data[self._st_max_index].left_up
        self._st_lo_interval = self.data[1].left_up.lo - self.data[0].left_up.lo

        logging.info('standard left up position: ' + str(self._st_min_pos))
        logging.info('standard right_down position: ' + str(self._st_max_pos))
        for i in range(1, len(self.data)):
            if (self.data[i].left_up.la != self.data[i - 1].left_up.la):
                self._st_la_interval = self.data[i].left_up.la - self.data[i - 1].left_up.la
                self._st_la_interval = - self._st_la_interval
                break

        self._st_lo_num = int(round((self._st_max_pos.lo - self._st_min_pos.lo) / self._st_lo_interval)) + 1
        self._st_la_num = int(round((self._st_min_pos.la - self._st_max_pos.la) / self._st_la_interval)) + 1

        """确保范围地图完全, 整齐"""
        count = 0
        prev = None
        for i in range(self._st_min_index, self._st_max_index + 1):
            if self.data[i].left_up.lo >= self._st_min_pos.lo and self.data[i].left_up.lo <= self._st_max_pos.lo:
                count += 1
                if prev is not None and self.data[i].left_up.la == prev.la:
                    assert(self.data[i].left_up.lo == prev.lo + self._st_lo_interval)
                elif prev is not None:
                    assert(self.data[i].left_up.la == prev.la - self._st_la_interval)
                prev = self.data[i].left_up

        assert(count == self._st_lo_num * self._st_la_num)
        logging.info('Picture number in this region: ' + str(count))

        self._st_fold = 1
        self._map_built = False
        self._width = self.data[0].width + 1   # 原来的width、height是255，与实际差1
        self._height = self.data[0].height + 1

    def check_map(self):
        """仔细检查map的完整性"""
        assert(self._map_built is True)
        count = 0
        for line in self.map:
            for record in line:
                count += 1
        for i in range(len(self.map)):
            for j in range(len(self.map[i])):
                if j != 0:
                    assert(abs((self.map[i][j].left_up.lo - self.map[i][j - 1].left_up.lo) - self._st_lo_interval / self._st_fold) <= 0.0000001)
                if i != 0:
                    assert(abs((self.map[i - 1][j].left_up.la - self.map[i][j].left_up.la) - self._st_la_interval) <= 0.0000001)
        assert(count == self._st_lo_num * self._st_la_num + (self._st_fold - 1) * (self._st_lo_num - 1) * self._st_la_num)

    def fold_image(self, pix_data1, pix_data2, index, fold):
        """fold是总的数目，index表示是第几个（从1开始，到fold-1）"""
        w_start = int(float(self._width) / float(fold) * index)
        w_end = int(float(self._width) / float(fold) * index)
        new_image = Image.new("RGB", (self._width - w_start + w_end + 1, self._height))
        pix = new_image.load()
        for i in range(self._width - w_start):
            for j in range(self._height):
                pix[i, j] = pix_data1[w_start + i, j]
        for i in range(w_end + 1):
            for j in range(self._height):
                pix[i + self._width - w_start, j] = pix_data2[i, j]
        new_image = new_image.resize((self._width, self._height))
        return new_image

    def create_folded_img(self, folded_file_dir, folded_index_file, fold):
        """创建图片文件，同时build map"""
        assert(fold > 1)
        self._st_fold = fold
        assert(self._map_built is False)
        index = []
        self.map = []
        for i in range(self._st_min_index, self._st_max_index + 1):
            if self.data[i].left_up.lo >= self._st_min_pos.lo and self.data[i].left_up.lo <= self._st_max_pos.lo:
                index.append(i)  # index记录下下标

        with open(folded_index_file, 'w') as f:
            # for check
            f.write(str(self._st_min_pos) + '\n')
            f.write(str(self._st_max_pos) + '\n')
            f.write(str(self._st_fold) + '\n')

            # do folding
            folded_img_count = 0
            for i in range(self._st_la_num):
                tmp = []
                for j in range(self._st_lo_num - 1):  # 最后一个不能fold
                    image_info1 = self.data[index[i * self._st_lo_num + j]]
                    image_info2 = self.data[index[i * self._st_lo_num + j + 1]]
                    tmp.append(self.data[index[i * self._st_lo_num + j]])
                    pix1 = Image.open(image_info1.file_path).load()
                    pix2 = Image.open(image_info2.file_path).load()
                    for k in range(1, fold):
                        folded_img_count += 1
                        file_path = os.path.join(folded_file_dir, '%d.jpg' % folded_img_count)
                        folded_img = self.fold_image(pix1, pix2, k, fold)
                        folded_img.save(file_path)
                        left_up_lo = self._st_lo_interval / fold * k + image_info1.left_up.lo
                        right_down_lo = left_up_lo + self._st_lo_interval
                        left_up_la = image_info1.left_up.la
                        right_down_la = image_info1.right_down.la
                        f.write(file_path + ' ' + str(left_up_lo) + ',' + str(left_up_la) + ' ' + str(right_down_lo) + ',' + str(right_down_la) + ' ' + str(self._width) + ' ' + str(self._height) + '\n')
                        tmp.append(ImageInfo(
                            file_path,
                            Position(lo=left_up_lo, la=left_up_la),
                            Position(lo=right_down_lo, la=right_down_la),
                            self._width,
                            self._height
                        ))
                # 补上最后一个
                tmp.append(self.data[index[i * self._st_lo_num + self._st_lo_num - 1]])
                self.map.append(tmp)
                logging.info(str(i + 1) + ' / ' + str(self._st_la_num))

        self._map_built = True
        self.check_map()

    def build_map(self):
        """map[_st_la_num][_st_lo_num]以映射的形式存放数据便于查找"""
        """如果有fold，则是map[_st_la_num][_st_lo_num*fold]"""
        assert(self._st_fold == 1)
        assert(self._map_built is False)
        self.map = []
        tmp = []
        index = []
        for i in range(self._st_min_index, self._st_max_index + 1):
            if self.data[i].left_up.lo >= self._st_min_pos.lo and self.data[i].left_up.lo <= self._st_max_pos.lo:
                index.append(i)  # index记录下下标
        for i in range(self._st_la_num):
            tmp = []
            for j in range(self._st_lo_num):
                tmp.append(self.data[index[i * self._st_lo_num + j]])
            self.map.append(tmp)

        self._map_built = True
        self.check_map()

    def build_map_with_exist_fold(self, folded_index_file, fold):
        """从文件中读入fold索引，从而构建合适的map"""
        assert(self._map_built is False)
        self.map = []
        tmp = []
        index = []
        for i in range(self._st_min_index, self._st_max_index + 1):
            if self.data[i].left_up.lo >= self._st_min_pos.lo and self.data[i].left_up.lo <= self._st_max_pos.lo:
                index.append(i)  # index记录下下标

        with open(folded_index_file, 'r') as f:
            # check
            assert(f.readline().strip() == str(self._st_min_pos))
            assert(f.readline().strip() == str(self._st_max_pos))
            assert(int(f.readline()) == fold)
            self._st_fold = fold

            # produce map
            count = 0
            for i in range(self._st_la_num):
                tmp = []
                for j in range(self._st_lo_num - 1):
                    tmp.append(self.data[index[i * self._st_lo_num + j]])
                    count += 1
                    for k in range(1, fold):
                        line = f.readline()
                        line = line.split(' ')
                        tmp.append(ImageInfo(line[0],
                                             Position(float(line[1].split(',')[0]), float(line[1].split(',')[1])),
                                             Position(float(line[2].split(',')[0]), float(line[2].split(',')[1])),
                                             int(line[3]),
                                             int(line[4]),
                                             ))
                        count += 1
                # 补上最后一个
                tmp.append(self.data[index[i * self._st_lo_num + self._st_lo_num - 1]])
                count += 1
                self.map.append(tmp)
        self._map_built = True
        self.check_map()
        logging.info('After adding the folded images, we have %d images in all.' % count)

    def get(self, position):
        """根据位置找图片"""
        """如果有fold只会找不是fold的那些图片"""
        assert(position >= self._st_min_pos and position <= self._st_max_pos)
        assert(self._map_built is True)
        i = int((self._st_min_pos.la - position.la) / self._st_la_interval)
        j = int((position.lo - self._st_min_pos.lo) / self._st_lo_interval) * self._st_fold
        return self.map[i][j]

    def __decide_class(self, num, scale):
        assert num >= scale[0]
        for i in range(1, len(scale)):
            if num >= scale[i - 1] and num < scale[i]:
                return i - 1
        return len(scale) - 1

    def ensemble_y_classification(self, y_file, scale=[0, 1]):
        """给map装上训练目标"""
        """scale [a,b,c,d] 1. >=a <b 2. >=b <c 3. >=c <d 4. >=d"""
        assert(self._map_built is True)
        for line in self.map:
            for record in line:
                record.y = 0

        not_in_region = 0
        with open(y_file) as f:
            for line in f:
                line = line.split(' ')
                la = float(line[-2])
                lo = float(line[-3])
                if not(Position(la=la, lo=lo).in_region(self._st_min_pos, self._st_max_pos)):
                    not_in_region += 1
                    continue
                i = int((self._st_min_pos.la - la) / self._st_la_interval)
                j = int((lo - self._st_min_pos.lo) / self._st_lo_interval) * self._st_fold
                for index in range(j, j + self._st_fold):
                    if index < self._st_lo_num:  # 最后一个没有fold要注意
                        if Position(la=la, lo=lo).in_region(self.map[i][index].left_up, self.map[i][index].right_down):
                            #  self.map[i][index].y += 1  # 商户数
                            try:
                                self.map[i][index].y += int(line[-6]) + 1
                            except Exception:
                                logging.warning('Can\'t concert to int: %s' % line[-6])
        if not_in_region != 0:
            logging.warning('%d line not in this region!' % not_in_region)

        for line in self.map:
            for i in range(len(line)):
                line[i].y = self.__decide_class(line[i].y, scale)

    def __build_file_records(self, balance_factor, balance_label):
        assert(self._map_built is True)
        count = 0
        for line in self.map:
            for record in line:
                if record.y == balance_label:
                    count += 1
        tmp_all_records = []
        for line in self.map:
            for record in line:
                r = FileRecord(record.file_path, record.y)
                tmp_all_records.append(r)
        random.shuffle(tmp_all_records)

        # do balance
        count_for_balance = 0
        all_records = []
        max_samples_balanced = int(count * balance_factor)
        if max_samples_balanced < count:
            logging.warning('Balance for label %d. Original number: %d. Balanced number: %d. ' % (balance_label, count, max_samples_balanced))
        for record in tmp_all_records:
            if record.y == balance_label:
                if count_for_balance >= max_samples_balanced:
                    continue
                count_for_balance += 1
            all_records.append(record)
        random.shuffle(all_records)

        record_y = {}
        for record in all_records:
            if record.y not in record_y:
                record_y[record.y] = 1
            else:
                record_y[record.y] += 1
        logging.warning('The number of each class: ' + str(record_y))

        return all_records

    def create_mxnet_index(self, out_file, balance_factor=1.0, balance_label=0):
        all_records = self.__build_file_records(balance_factor, balance_label)
        with open(out_file, 'w') as f:
            for i in range(len(all_records)):
                f.write('%d %d %s\n' % (i, all_records[i].y, all_records[i].file_path))

    def create_tensorflow_imgs(self, out_dir, balance_factor=1.0, balance_label=0, train_factor=1.0):
        all_records = self.__build_file_records(balance_factor, balance_label)

        # build dir
        record_y = {}
        for record in all_records:
            if record.y not in record_y:
                record_y[record.y] = 1
        if os.path.exists(out_dir):
            shutil.rmtree(out_dir)
        os.makedirs(out_dir)
        train_dir = os.path.join(out_dir, 'train')
        validation_dir = os.path.join(out_dir, 'validation')
        os.makedirs(train_dir)
        os.makedirs(validation_dir)
        for y in record_y:
            os.makedirs(os.path.join(train_dir, str(y)))
            os.makedirs(os.path.join(validation_dir, str(y)))

        train_num = int(len(all_records) * train_factor)
        train_records = all_records[0:train_num]
        validation_records = all_records[train_num:len(all_records)]

        # for training
        for record in train_records:
            path_name = os.path.join(train_dir, str(record.y))
            new_file_path = os.path.join(path_name, os.path.basename(record.file_path))
            shutil.copy(record.file_path, new_file_path)

        # for validation
        for record in validation_records:
            path_name = os.path.join(validation_dir, str(record.y))
            new_file_path = os.path.join(path_name, os.path.basename(record.file_path))
            shutil.copy(record.file_path, new_file_path)

        with open(os.path.join(out_dir, 'label.txt') ,'w') as f:
            for label in record_y:
                f.write(str(label) + '\n')


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    i = ImageData('/home/hzy/project/dp_up/data/beijing/L17.txt',
                  Position(la=40.0756948689, lo=116.1727796185),
                  Position(la=39.7421689280, lo=116.5992171065))
    # i.create_folded_img('/home/hzy/project/dp_up/data/beijing/L17_append', 'L17_append.txt', 3)
    i.build_map_with_exist_fold('/home/hzy/project/dp_up/data/beijing/L17_append.txt', 3)
    # i.build_map()
    # print(i.get(Position(la=40.0135332578, lo=116.3840892949)))
    i.ensemble_y_classification('/home/hzy/project/dp_up/data/restaurant/Beijing/beijing_wgs.txt')
    i.create_mxnet_index('mxnet_index_beijing.txt')
